// Copyright (c) 2020 Presto Labs Pte. Ltd.
// Author: donggu

static const char USAGE[] = R"(
    Order Plotter with optional Feed.

    Usage:
      plot --driver=<config_path>
      plot --mea=<mea> --product=<norm_symbol>
           --feed_machine=<feed_machine> --recipe=<recipe>
           --account=<account> --strat_machine=<strat_machine>
           [--start=<start>] [--duration=<duration>] [--end=<end>]

    Options:
      -h --help             Show this screen.
      --start=<start>       YYYY-MM-DD
      --end=<end>           YYYY-MM-DD
      --duration=<duration> PT1H
)";

#include <string>
#include <tuple>

#include <boost/algorithm/string.hpp>
#include <docopt/docopt.h>

#include "coin/proto/coin_order_enums.pb.h"
#include "coin/proto/coin_order_gateway.pb.h"
#include "coin2/app/driver.h"
#include "coin2/base/config/app_config.h"
#include "coin2/base/config/args.h"
#include "coin2/exchange/base/order/order.h"
#include "coin2/strategy/strategy.h"
#include "vitro.h"  // NOLINT

using coin::proto::OrderDirection;
using coin::proto::OrderEvent;
using coin::proto::OrderSide;
using coin2::exchange::base::feed::FeedUpdate;
using coin2::exchange::base::order::OrderUpdate;
using namespace coin2::app;
using namespace coin2::base::config;

int BUY = 0;
int SELL = 1;

struct AreaSeries {
  std::vector<int64_t> x;
  std::vector<double> y1;
  std::vector<double> y2;
};

struct Series {
  std::vector<int64_t> x;
  std::vector<double> y;
};

int trade_side(const Trade& trade) {
  if (trade.side == coin2::exchange::feed::TRADE_BUY_SIDE) {
    return 0;
  }
  if (trade.side == coin2::exchange::feed::TRADE_SELL_SIDE) {
    return 1;
  }
  throw std::runtime_error("unknown trade side");
}

std::map<std::string, AreaSeries> bbo;
std::map<std::string, Series> trade;
std::map<std::string, Series[2]> submit;
std::map<std::string, Series[2]> fill;
std::map<std::string, Series[2]> cancel;
std::map<std::string, Series[2]> reject;
std::map<std::string, Series> balance;
std::map<std::string, Series> position;
std::map<std::string, Series> unknown;

int order_side(OrderSide side) {
  switch (side) {
    case OrderSide::BUY_ORDER:
    case OrderSide::BUY_OPEN_ORDER:
    case OrderSide::BUY_CLOSE_ORDER:
      // case OrderSide::BUY_CLOSE_CARRYOVER_ORDER:
      return 0;
    case OrderSide::SELL_ORDER:
    case OrderSide::SELL_SHORT_ORDER:
    case OrderSide::SELL_OPEN_ORDER:
    case OrderSide::SELL_CLOSE_ORDER:
      // case OrderSide::SELL_CLOSE_CARRYOVER_ORDER:
      return 1;
    default:
      NOTREACHED() << OrderSide_Name(side);
      return OrderDirection::DIRECTION_UNINITIALIZED;
  }
}

struct Strategy : public ::coin2::strategy::IStrategy {
  void onTradeFeed(const FeedUpdate& upd) override {
    auto id = upd.system_name() + ":" + upd.product().absolute_norm();
    auto& series = trade[id];
    series.x.push_back(upd.timestamp());
    series.y.push_back(upd.trade().price);
  }
  void onBookFeed(const FeedUpdate& upd) override {
    auto& book = upd.book();
    if (book.Ask0() && book.Bid0()) {
      auto id = upd.system_name() + ":" + upd.product().absolute_norm();
      auto& series = bbo[id];
      series.x.push_back(upd.timestamp());
      series.y1.push_back(upd.book().Bid0()->price);
      series.y2.push_back(upd.book().Ask0()->price);
    }
  }

  void onAccountInfo(const OrderUpdate& upd) override {
    for (auto& [product, pos] : upd.account_info().GetPositions()) {
      auto key = upd.system_name() + ":" + product;
      auto& series = position[key];
      series.x.push_back(upd.timestamp());
      series.y.push_back(pos.net_position());
    }
    for (auto& [currency, b] : upd.account_info().GetBalances()) {
      auto key = upd.system_name() + ":" + currency;
      auto& series = balance[key];
      series.x.push_back(upd.timestamp());
      series.y.push_back(b.total());
    }
  }
  void onAccountOrder(const OrderUpdate& upd) override {
    auto& ev = upd.product_order_info().event();
    if (ev.type() == OrderEvent::ORDER_SUBMITTED) {
      auto key = upd.system_name() + ":" + ev.symbol();
      auto& series = submit[key][order_side(ev.order_side())];
      series.x.push_back(upd.timestamp());
      series.y.push_back(ev.order_price());
    } else if (ev.type() == OrderEvent::ORDER_FILLED) {
      auto key = upd.system_name() + ":" + ev.symbol();
      auto& series = fill[key][order_side(ev.order_side())];
      series.x.push_back(upd.timestamp());
      series.y.push_back(ev.order_price());
    } else if (ev.type() == OrderEvent::ORDER_REJECTED || ev.type() == OrderEvent::ORDER_ERROR) {
      auto key = upd.system_name() + ":" + ev.symbol();
      auto& series = reject[key][order_side(ev.order_side())];
      series.x.push_back(upd.timestamp());
      series.y.push_back(ev.order_price());
    } else if (ev.type() == OrderEvent::CANCEL_CONFIRMED) {
      auto key = upd.system_name() + ":" + ev.symbol();
      auto& series = cancel[key][order_side(ev.order_side())];
      series.x.push_back(upd.timestamp());
      series.y.push_back(ev.order_price());
    } else {
      /*auto key = upd.system_name() + ":" + ev.symbol();
      auto& series = unknown[key];
      series.x.push_back(upd.timestamp());
      series.y.push_back(ev.order_price());*/
    }
  }
};

DriverConfig parse_cli(int argc, char* argv[]) {
  std::map<std::string, docopt::value> args = docopt::docopt(
      USAGE,
      {argv + 1, argv + argc},
      true,                 // show help if requested
      "Coin2 Order Plot");  // version string

  printf("!!!\n");

  coin2::app::DriverConfig driver;
  if (args["--driver"]) {
    driver.MergeFrom(coin2::base::config::DriverConfigFromJson(args["--driver"].asString()));
  } else {
    // feed
    {
      auto& exchanges = *driver.mutable_feed()->mutable_exchanges();
      auto& exchange = exchanges[args["--mea"].asString()];

      exchange.set_mea(args["--mea"].asString());
      exchange.set_recipe(args["--recipe"].asString());
      exchange.set_type(coin2::exchange::feed::FeedSystemType::FEED_SYSTEM_TYPE_ARCHIVE);
      exchange.mutable_products()->add_norms(args["--product"].asString());
      exchange.mutable_archive()->set_archive_type(coin2::exchange::feed::ArchiveConfig::FASTFEED);
      exchange.mutable_archive()->set_machine(args["--feed_machine"].asString());
    }
    // order
    {
      auto& exchanges = *driver.mutable_order()->mutable_exchanges();
      auto& exchange = exchanges[args["--mea"].asString()];

      exchange.set_mea(args["--mea"].asString());
      exchange.set_type(coin2::exchange::order::OrderSystemType::ORDER_SYSTEM_TYPE_ARCHIVE);
      exchange.mutable_products()->add_norms(args["--product"].asString());
      exchange.mutable_archive()->set_machine(args["--strat_machine"].asString());
      exchange.mutable_archive()->set_owner(args["--account"].asString());
    }
  }

  {
    auto& feed_interval =
        *driver.mutable_feed()->mutable_common()->mutable_archive()->mutable_interval();
    auto& order_interval =
        *driver.mutable_order()->mutable_common()->mutable_archive()->mutable_interval();
    if (args["--start"]) {
      feed_interval.set_start(args["--start"].asString());
      order_interval.set_start(args["--start"].asString());
    }
    if (args["--duration"]) {
      feed_interval.set_duration(args["--duration"].asString());
      order_interval.set_duration(args["--duration"].asString());
    }
    if (args["--end"]) {
      feed_interval.set_duration(args["--end"].asString());
      order_interval.set_duration(args["--end"].asString());
    }
  }

  return driver;
}

int main(int argc, char* argv[]) {
  auto app = parse_cli(argc, argv);

  /* ----------------------------------- run ---------------------------------- */
  printf("reading ..\n");

  Strategy strategy;
  ::coin2::app::Driver driver(app, &strategy);
  driver.Run();

  /* ---------------------------------- plot ---------------------------------- */
  printf("Plotting ..\n");

  auto out = std::experimental::filesystem::path("/tmp/order_plot.png");
  double marker_size = 0;
  if (out.extension() == ".png") {
    marker_size = 10.0;
  } else if (out.extension() == ".svg") {
    marker_size = 2.0;
  } else {
    throw std::runtime_error("only supports .png or .svg");
  }

  Figure fig(3, 1);  // price(feed, order), pos, balance
  fig.title = "Order Replay";
  fig.width = 1600;
  fig.height = 1200;
  fig.height_ratios = {10, 3, 3};

  /* ---------------------------- plot 1/3 - price ---------------------------- */
  {
    auto& ax = fig.axes(1, 1);
    ax.is_x_nanotimestamps = true;
    std::string text;

    // feed
    for (auto& [key, series] : trade) {
      auto& line = ax.line(key + ":trade", series.x, series.y);
      line.alpha = 0.4;
      line.width = 0.2;
      line.color = "blue";
      line.drawstyle = "steps-post";
      text += key + ":trade = " + std::to_string(series.x.size()) + '\n';
    }
    for (auto& [key, series] : bbo) {
      auto& area = ax.area(key + ":bbo", series.x, series.y1, series.y2);
      area.alpha = 0.6;
      area.color = "khaki";
      text += key + ":bbo = " + std::to_string(series.x.size()) + '\n';
    }
    // order
    for (auto& [key, series] : submit) {
      {
        auto& scatter = ax.scatter(key + ":submit-buy", series[0].x, series[0].y);
        scatter.marker_type = 'o';
        scatter.marker_size = marker_size;
        scatter.marker_face_color = "none";
        scatter.width = 0.1;
        scatter.alpha = 0.8;
        scatter.color = "g";
      }
      {
        auto& scatter = ax.scatter(key + ":submit-sell", series[1].x, series[1].y);
        scatter.marker_type = 'o';
        scatter.marker_size = marker_size;
        scatter.marker_face_color = "none";
        scatter.width = 0.1;
        scatter.alpha = 0.8;
        scatter.color = "r";
      }
      text += key + ":submit-buy = " + std::to_string(series[0].x.size()) + '\n';
      text += key + ":submit-sell = " + std::to_string(series[1].x.size()) + '\n';
    }
    for (auto& [key, series] : fill) {
      {
        auto& scatter = ax.scatter(key + ":fill-buy", series[0].x, series[0].y);
        scatter.marker_type = 'o';
        scatter.marker_size = marker_size / 2.0;
        scatter.width = 0.1;
        scatter.alpha = 0.6;
        scatter.color = "g";
      }
      {
        auto& scatter = ax.scatter(key + ":fill-sell", series[1].x, series[1].y);
        scatter.marker_type = 'o';
        scatter.marker_size = marker_size / 2.0;
        scatter.width = 0.1;
        scatter.alpha = 0.6;
        scatter.color = "r";
      }
      text += key + ":fill-buy = " + std::to_string(series[0].x.size()) + '\n';
      text += key + ":fill-sell = " + std::to_string(series[1].x.size()) + '\n';
    }
    for (auto& [key, series] : cancel) {
      {
        auto& scatter = ax.scatter(key + ":cancel-buy", series[0].x, series[0].y);
        scatter.marker_type = '4';
        scatter.marker_size = marker_size;
        scatter.width = 0.1;
        scatter.alpha = 0.4;
        scatter.color = "g";
      }
      {
        auto& scatter = ax.scatter(key + ":cancel-sell", series[1].x, series[1].y);
        scatter.marker_type = '4';
        scatter.marker_size = marker_size;
        scatter.width = 0.1;
        scatter.alpha = 0.4;
        scatter.color = "r";
      }
      text += key + ":cancel-buy = " + std::to_string(series[0].x.size()) + '\n';
      text += key + ":cancel-sell = " + std::to_string(series[1].x.size()) + '\n';
    }
    for (auto& [key, series] : reject) {
      {
        auto& scatter = ax.scatter(key + ":reject-buy", series[0].x, series[0].y);
        scatter.marker_type = 'x';
        scatter.marker_size = marker_size;
        scatter.width = 0.1;
        scatter.alpha = 0.4;
        scatter.color = "g";
      }
      {
        auto& scatter = ax.scatter(key + ":reject-sell", series[1].x, series[1].y);
        scatter.marker_type = 'x';
        scatter.marker_size = marker_size;
        scatter.alpha = 0.4;
        scatter.color = "r";
      }
      text += key + ":reject-buy = " + std::to_string(series[0].x.size()) + '\n';
      text += key + ":reject-sell = " + std::to_string(series[1].x.size()) + '\n';
    }
    /*for (auto& [key, series] : unknown) {
      {
        auto& scatter = ax.scatter(key + ":unknown", series.x, series.y);
        scatter.marker_type = "$?$";
        scatter.color = "g";
      }
      text += key + ":unknown = " + std::to_string(series.x.size()) + '\n';
    }*/

    boost::trim_right(text);
    ax.text(text, 0.95, 0.05);
  }

  /* ---------------------------- plot 2/3 - pos ---------------------------- */
  {
    auto& ax = fig.axes(2, 1);
    ax.title = "position";
    ax.is_x_nanotimestamps = true;
    std::string text;

    // feed
    for (auto& [key, series] : position) {
      auto& line = ax.line(key + ":position", series.x, series.y);
      line.alpha = 1.0;
      line.width = 0.2;
      line.drawstyle = "steps-post";
      text += key + ":position = " + std::to_string(series.x.size()) + '\n';
    }

    boost::trim_right(text);
    ax.text(text, 0.95, 0.05);
  }

  /* --------------------------- plot 3/3 - balance --------------------------- */
  {
    auto& ax = fig.axes(3, 1);
    ax.title = "balance";
    ax.is_x_nanotimestamps = true;
    std::string text;

    // feed
    for (auto& [key, series] : balance) {
      auto& line = ax.line(key + ":balance", series.x, series.y);
      line.alpha = 1.0;
      line.width = 0.2;
      line.color = "black";
      line.drawstyle = "steps-post";
      text += key + ":balance = " + std::to_string(series.x.size()) + '\n';
    }

    boost::trim_right(text);
    ax.text(text, 0.95, 0.05);
  }

  Matplot matplot(fig);
  matplot.save(out);
  return 0;
}
